import random
import time
import numpy as np
import Processor.Converter as Converter
from PyQt5.QtWidgets import QDialog
from Processor.noise.gasuss_design import Ui_GasussDialog
from Processor.noise.sp_design import Ui_SPDialog


class GasussDialog(QDialog, Ui_GasussDialog):
    def __init__(self, parent):
        super().__init__(parent)
        self.setupUi(self)
        self.__initUI(parent)

    def __initUI(self, parent):
        self.setWindowTitle('高斯噪声')
        # 定位
        self.resize(180, 170)
        pos = parent.geometry()
        self.move((pos.width() - self.width()) / 2, (pos.height() - self.height()) / 2)
        # 信号
        #       对话框确定
        self.DialogButtonBox.accepted.connect(self.accept)
        self.DialogButtonBox.rejected.connect(self.reject)


class SPDialog(QDialog, Ui_SPDialog):
    def __init__(self, parent):
        super().__init__(parent)
        self.setupUi(self)
        self.__initUI(parent)

    def __initUI(self, parent):
        self.setWindowTitle('椒盐噪声')
        # 定位
        self.resize(180, 170)
        pos = parent.geometry()
        self.move((pos.width() - self.width()) / 2, (pos.height() - self.height()) / 2)
        # 信号
        #       对话框确定
        self.DialogButtonBox.accepted.connect(self.accept)
        self.DialogButtonBox.rejected.connect(self.reject)


def spNoise(image, prob):
    """
    添加椒盐噪声
    prob:噪声比例
    :param image:
    :param prob:
    :return:
    """
    row, col, _ = image.shape
    output = image.copy()
    thres = 1 - prob
    for i in range(row):
        for j in range(col):
            rdn = random.random()
            if rdn < prob:
                output[i][j] = 0
            elif rdn > thres:
                output[i][j] = 255
    return output


def gasussNoise(image, mean=0, var=0.001):
    """
    :param image: 数据
    :param mean: 均值
    :param var: 方差
    :return:
    """
    image = np.array(image/255, dtype=float)
    noise = np.random.normal(mean, var ** 0.5, image.shape)
    out = image + noise
    if out.min() < 0:
        low_clip = -1.
    else:
        low_clip = 0.
    out = np.clip(out, low_clip, 1.0)
    out = np.uint8(out*255)
    return out


def newGasussNoise(window):
    dialog = GasussDialog(window)
    imageWindow = window.currentImageWindow
    if not imageWindow:
        return
    layer = imageWindow.currentLayer
    if not layer:
        return
    if dialog.exec_():
        start = time.process_time()
        mean = dialog.doubleSpinBox_mean.value()
        var = dialog.doubleSpinBox_var.value()
        image = layer.image
        mat = Converter.QImage2Numpy(image)
        mat[..., 0:3] = gasussNoise(mat[..., 0:3], mean, var)
        layer.updatePixmap(Converter.Numpy2QImage(mat))
        end = time.process_time()
        print('GasussNoise 耗时:', end - start)


def newSPNoise(window):
    imageWindow = window.currentImageWindow
    if not imageWindow:
        return
    layer = imageWindow.currentLayer
    if not layer:
        return
    dialog = SPDialog(window)
    if dialog.exec_():
        start = time.process_time()
        prob = dialog.doubleSpinBox_S.value()
        image = layer.image
        mat = Converter.QImage2Numpy(image)
        mat[..., 0:3] = spNoise(mat[..., 0:3], prob)
        layer.updatePixmap(Converter.Numpy2QImage(mat))
        end = time.process_time()
        print('SPNoise 耗时:', end - start)